from itertools import combinations
import matplotlib.pyplot as plt

def load_dataset():
    """模拟超市购物篮数据集"""
    return [
        ['牛奶', '面包', '黄油'],
        ['啤酒', '面包'],
        ['牛奶', '啤酒', '面包', '鸡蛋'],
        ['面包', '黄油'],
        ['牛奶', '面包', '鸡蛋'],
        ['啤酒', '鸡蛋'],
        ['牛奶', '黄油'],
        ['牛奶', '面包', '黄油', '鸡蛋']
    ]

def create_candidates(itemset, k):
    """生成候选k-项集"""
    candidates = set()
    items = list(itemset)
    for i in range(len(items)):
        for j in range(i+1, len(items)):
            union_set = items[i].union(items[j])
            if len(union_set) == k:
                candidates.add(union_set)
    return candidates

def calculate_support(itemset, transactions):
    """计算项集的支持度"""
    count = 0
    for transaction in transactions:
        if itemset.issubset(transaction):
            count += 1
    return count / len(transactions)

def apriori(dataset, min_support=0.3):
    """Apriori算法实现"""
    transactions = list(map(set, dataset))
    all_items = set([item for transaction in transactions for item in transaction])
    
    supports = {}
    k = 1
    current_items = {frozenset([item]) for item in all_items}
    
    print("挖掘频繁项集过程:")
    print("-" * 40)
    
    while current_items:
        current_supports = {}
        print(f"\n{k}-项集候选:")
        
        for item in current_items:
            support = calculate_support(item, transactions)
            if support >= min_support:
                current_supports[item] = support
                print(f"  {set(item)}: {support:.2f}")
        
        supports.update(current_supports)
        current_items = create_candidates(current_supports.keys(), k + 1)
        k += 1
    
    return supports

def generate_association_rules(supports, min_confidence=0.7):
    """生成关联规则"""
    print(f"\n关联规则 (最小置信度 {min_confidence}):")
    print("-" * 40)
    
    rules = []
    for itemset in supports:
        if len(itemset) > 1:
            itemset_list = list(itemset)
            # 生成所有可能的规则
            for i in range(1, len(itemset_list)):
                for antecedent in combinations(itemset_list, i):
                    antecedent = frozenset(antecedent)
                    consequent = itemset - antecedent
                    
                    if antecedent in supports and consequent:
                        confidence = supports[itemset] / supports[antecedent]
                        if confidence >= min_confidence:
                            rules.append((antecedent, consequent, confidence))
                            print(f"  {set(antecedent)} => {set(consequent)} : 置信度={confidence:.2f}")
    
    return rules

def visualize_results(supports):
    """可视化频繁项集支持度"""
    itemsets = [f"L{len(k)}" for k in supports.keys()]
    support_values = list(supports.values())
    
    plt.figure(figsize=(12, 6))
    bars = plt.bar(range(len(support_values)), support_values, color='skyblue', alpha=0.7)
    plt.xlabel('频繁项集')
    plt.ylabel('支持度')
    plt.title('频繁项集支持度分布')
    plt.xticks(range(len(itemsets)), itemsets, rotation=45)
    
    # 在柱子上显示数值
    for bar, value in zip(bars, support_values):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01, 
                f'{value:.2f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('apriori_results.png', dpi=300, bbox_inches='tight')
    plt.show()

def run_apriori():
    print("=" * 50)
    print("Apriori算法实验")
    print("=" * 50)
    
    # 加载数据集
    dataset = load_dataset()
    print("数据集:")
    for i, transaction in enumerate(dataset, 1):
        print(f"  交易{i}: {transaction}")
    
    # 运行Apriori算法
    min_support = 0.3
    supports = apriori(dataset, min_support=min_support)
    
    print(f"\n所有频繁项集 (最小支持度 {min_support}):")
    print("-" * 40)
    for itemset, support in sorted(supports.items(), key=lambda x: (-len(x[0]), -x[1])):
        print(f"  {set(itemset)}: {support:.2f}")
    
    # 生成关联规则
    rules = generate_association_rules(supports, min_confidence=0.7)
    
    # 可视化结果
    visualize_results(supports)
    
    return len(supports), len(rules)

if __name__ == "__main__":
    run_apriori()